package han.cirno.corrupt.components.cnn

import android.graphics.Bitmap
import han.cirno.corrupt.util.SerialUtil
import han.cirno.larva.model.LarvaModel
import han.cirno.larva.operation.OperationResolver
import han.cirno.larva.operation.Operations
import org.tensorflow.lite.Interpreter
import java.io.InputStream
import java.nio.ByteBuffer
import java.nio.FloatBuffer

class LarvaWrapper private constructor(
    private val interpreter:Interpreter,
    private val outputMap:Array<String>,
    private val operations: Operations){

    fun predict(bitmap: Bitmap):String{
        val larvaImage=LarvaAndroidUtil.fromBitmap(bitmap)
        val resolve = OperationResolver.resolve(larvaImage, operations)
        val stringBuilder=StringBuilder()
        for (image in resolve) {
            val inputVector = image.toVector()
            val output=Array(1){FloatArray(outputMap.size)}
            val floatBuffer = FloatBuffer.wrap(inputVector)
            interpreter.run(floatBuffer,output)
            stringBuilder.append(getString(output[0]))
        }
        return stringBuilder.toString()
    }

    private fun getString(vector:FloatArray):String{
        var maxIndex=-1
        var maxValue=-114514f
        for((i,v) in vector.withIndex()){
            if(maxValue<=v){
                maxValue=v
                maxIndex=i
            }
        }
        return outputMap[maxIndex]
    }

    companion object{
        @JvmStatic
        fun fromStream(stream: InputStream):LarvaWrapper?{
            val larvaModel= SerialUtil.deserialize<LarvaModel>(stream) ?: return null
            val tfLiteFileBytes = larvaModel.tfLiteFileBytes
            val byteBuffer:ByteBuffer=ByteBuffer.allocateDirect(tfLiteFileBytes.size)
            byteBuffer.put(tfLiteFileBytes)
            byteBuffer.flip()
            val larvaWrapper=LarvaWrapper(
                Interpreter(byteBuffer),
                larvaModel.outputMap,
                larvaModel.operations
            )
            return larvaWrapper
        }
    }
}